# ====== 全局扰动参数 ======
INTRINSIC_PERTURB_STD = 1000     # 内参扰动标准差（像素）
EULER_PERTURB_STD = 0.5         # 欧拉角扰动标准差（度）
TRANSLATION_PERTURB_STD = 50    # 平移扰动标准差（mm）
PIXEL_PERTURB_STD = 0         # 图像质心扰动标准差（像素）
TARGET_CALIB_ERR_STD = 0        # 靶标标定误差标准差（mm）
TS_MEASURE_ERR_STD = 0          # 全站仪测量误差标准差（mm）
# ==========================

import numpy as np
# 设置numpy打印选项，取消科学计数法
np.set_printoptions(suppress=True, precision=6)

from data import (
    read_csv_data, get_initial_camera_params,
    get_initial_camera_pose, get_target_points, get_object_points
)
import numpy as np
from algorithm import euler_to_matrix
from scipy.optimize import least_squares
import cv2
import tkinter as tk
from tkinter import messagebox

def show_iteration_dialog(Otc, Otc_estimated, iteration_num):
    """显示迭代对话框"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    
    message = f"迭代 {iteration_num}:\n\n"
    message += f"Otc: {Otc}\n"
    message += f"Otc_estimated: {Otc_estimated}\n"
    message += f"误差: {np.linalg.norm(Otc - Otc_estimated):.2f} mm"
    
    messagebox.showinfo("迭代信息", message)
    root.destroy()

def show_iteration_dialog_cost(cost, details, iteration_num):
    """显示迭代cost对话框"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    
    message = f"迭代 {iteration_num} 当前cost: {cost:.2f}\n\n" + details
    messagebox.showinfo("迭代cost信息", message)
    root.destroy()

def show_final_origin_errors(origin_errors, Otc_list, Otc_estimated_list, final_cost):
    """用tkinter的列表控件显示最终每个位置原点的几何一致性误差（各方向分量）、Otc、Otc_estimated，并显示最终cost和累加过程"""
    root = tk.Tk()
    root.title("最终原点误差")
    root.geometry("1000x700")
    frame = tk.Frame(root)
    frame.pack(fill=tk.BOTH, expand=True)
    scrollbar = tk.Scrollbar(frame)
    scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
    listbox = tk.Listbox(frame, width=140, height=35, font=("Consolas", 10))
    listbox.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
    listbox.config(yscrollcommand=scrollbar.set)
    scrollbar.config(command=listbox.yview)

    listbox.insert(tk.END, f"最终优化参数下，各位置原点几何一致性误差（单位:mm）")
    listbox.insert(tk.END, f"最终cost: {final_cost:.3f}")
    listbox.insert(tk.END, "")
    listbox.insert(tk.END, f"每个位姿误差平方及累加和：")
    sum_sq = 0.0
    for i, (err, Otc, Otc_est) in enumerate(zip(origin_errors, Otc_list, Otc_estimated_list)):
        err_sq = np.sum(np.array(err) ** 2)
        sum_sq += err_sq
        listbox.insert(tk.END, f"靶标{i+1}: 误差平方={err_sq:.3f}，累加和={sum_sq:.3f}")
    listbox.insert(tk.END, "")
    for i, (err, Otc, Otc_est) in enumerate(zip(origin_errors, Otc_list, Otc_estimated_list)):
        listbox.insert(tk.END, f"靶标{i+1}:")
        listbox.insert(tk.END, f"  Otc:          [{Otc[0]:.3f}, {Otc[1]:.3f}, {Otc[2]:.3f}]")
        listbox.insert(tk.END, f"  Otc_estimated:[{Otc_est[0]:.3f}, {Otc_est[1]:.3f}, {Otc_est[2]:.3f}]")
        listbox.insert(tk.END, f"  error:        [x={err[0]:.3f}, y={err[1]:.3f}, z={err[2]:.3f}]")
        listbox.insert(tk.END, "")
    btn = tk.Button(root, text="关闭", command=root.destroy)
    btn.pack(pady=10)
    root.mainloop()

def residual_func(params, Otw_list, image_points_list, object_points_list):
    """计算残差函数，几何一致性误差用PnP解算T_tc"""
    fx, fy, cx, cy = params[:4]  # 相机内参
    euler = params[4:7]          # 相机欧拉角
    R_cw = euler_to_matrix(euler)
    R_wc = R_cw.T
    T_cw = params[7:10]
    T_wc = -R_wc @ T_cw

    res = []
    details = ""
    iteration_num = getattr(residual_func, 'iteration_num', 0) + 1
    residual_func.iteration_num = iteration_num
    total_origin = 0
    for i, (Otw, img_pts, obj_pts) in enumerate(zip(Otw_list, image_points_list, object_points_list)):
        # 构建相机内参矩阵
        K = np.array([[fx, 0, cx],
                      [0, fy, cy],
                      [0,  0,  1]])
        # PnP解算R_tc, T_tc
        obj_pts = np.array(obj_pts, dtype=np.float32)
        img_pts = np.array(img_pts, dtype=np.float32)
        success, rvec, tvec = cv2.solvePnP(obj_pts, img_pts, K, None, flags=cv2.SOLVEPNP_EPNP)
        if not success:
            # 若PnP失败，残差设为大值
            origin_error = np.array([1e6, 1e6, 1e6])
        else:
            R_tc, _ = cv2.Rodrigues(rvec)
            T_tc = tvec.flatten()
            # 计算Otc和Otc_estimated
            Otc = R_wc @ Otw + T_wc
            Ot = np.zeros(3)
            Otc_estimated = R_tc @ Ot + T_tc
            origin_error = Otc - Otc_estimated
        res.extend(origin_error)
        origin_norm = np.linalg.norm(origin_error)
        total_origin += origin_norm
        details += f"靶标{i+1}:\n  origin_error: {origin_error}\n  范数: {origin_norm:.2f} mm\n\n"
    cost = np.sum(np.array(res) ** 2)
    details += f"平均原点误差: {total_origin/len(Otw_list):.2f} mm\n"
    #show_iteration_dialog_cost(cost, details, iteration_num)
    return np.array(res)

def main():
    # 读取CSV文件
    file_path = '靶标数据_20250616_154424.csv'
    data = read_csv_data(file_path)
    
    # 获取初始相机参数
    initial_intrinsic = get_initial_camera_params(data)
    initial_extrinsic = get_initial_camera_pose(data)
    
    print("initial_intrinsic", initial_intrinsic)
    print("initial_extrinsic", initial_extrinsic)
    
    # ====== 添加扰动 ======
    np.random.seed(42)  # 保证可复现
    intrinsic_perturb = np.random.normal(0, INTRINSIC_PERTURB_STD, size=4)
    euler_perturb = np.random.normal(0, EULER_PERTURB_STD, size=3)
    translation_perturb = np.random.normal(0, TRANSLATION_PERTURB_STD, size=3)
    
    perturbed_intrinsic = initial_intrinsic + intrinsic_perturb
    # 确保cx, cy在图像范围内，并且更接近真值
    cx_true = initial_intrinsic[2]
    cy_true = initial_intrinsic[3]
    cx_range = 100  # 允许100像素的偏差
    cy_range = 100
    perturbed_intrinsic[2] = np.clip(perturbed_intrinsic[2], cx_true - cx_range, cx_true + cx_range)
    perturbed_intrinsic[3] = np.clip(perturbed_intrinsic[3], cy_true - cy_range, cy_true + cy_range)
    
    perturbed_extrinsic = initial_extrinsic.copy()
    perturbed_extrinsic[:3] += euler_perturb
    perturbed_extrinsic[3:] += translation_perturb
    # ======================
    
    # 获取靶标点数据
    Otw_list, _, image_points_list, _, _ = get_target_points(data)
    object_points = get_object_points()
    object_points_noisy = object_points + np.random.normal(0, TARGET_CALIB_ERR_STD, size=object_points.shape)
    object_points_list = [object_points_noisy for _ in range(len(Otw_list))]
    
    Otw_list_noisy = [Otw + np.random.normal(0, TS_MEASURE_ERR_STD, size=3) for Otw in Otw_list]
    
    # ====== 图像质心扰动 ======
    perturbed_image_points_list = []
    for img_pts in image_points_list:
        noise = np.random.normal(0, PIXEL_PERTURB_STD, size=img_pts.shape)
        perturbed_image_points_list.append(img_pts + noise)
    # =========================
    
    # 合并初始参数
    initial_guess = np.concatenate([perturbed_intrinsic, perturbed_extrinsic])
    
    # 计算初始cost
    initial_residuals = residual_func(initial_guess, Otw_list_noisy, perturbed_image_points_list, object_points_list)
    initial_cost = np.sum(initial_residuals ** 2)
    print("\n初始cost:", initial_cost)
    
    # 定义参数边界
    cx_true = initial_intrinsic[2]
    cy_true = initial_intrinsic[3]
    cx_range = 100  # 允许100像素的偏差
    cy_range = 100
    bounds = (
        [0, 0, cx_true - cx_range, cy_true - cy_range, -180, -180, -180, -np.inf, -np.inf, -np.inf],  # 下界
        [np.inf, np.inf, cx_true + cx_range, cy_true + cy_range, 180, 180, 180, np.inf, np.inf, np.inf]  # 上界
    )
    
    # 使用最小二乘法优化
    result = least_squares(
        residual_func,
        initial_guess,
        args=(Otw_list_noisy, perturbed_image_points_list, object_points_list),
        bounds=bounds,
        method='trf',
        verbose=1
    )
    
    # 解析优化结果
    optimized_params = result.x
    fx, fy, cx, cy = optimized_params[:4]
    euler = optimized_params[4:7]
    T_wc_opt = optimized_params[7:10]
    R_cw_opt = euler_to_matrix(euler)
    R_wc_opt = R_cw_opt.T
    T_cw_opt = T_wc_opt
    
    # 构建相机内参矩阵
    K = np.array([[fx, 0, cx],
                  [0, fy, cy],
                  [0,  0,  1]])
    
    print("optimized_params", optimized_params)
    print("K", K)
    print("R_wc_opt", R_wc_opt)
    print("T_wc_opt", T_wc_opt)

    # ====== 真值对比 ======
    print("\n=== 优化结果与真值对比 ===")
    print(f"fx: 优化={fx:.6f}, 真值={initial_intrinsic[0]:.6f}, 差值={fx-initial_intrinsic[0]:.6f}")
    print(f"fy: 优化={fy:.6f}, 真值={initial_intrinsic[1]:.6f}, 差值={fy-initial_intrinsic[1]:.6f}")
    print(f"cx: 优化={cx:.6f}, 真值={initial_intrinsic[2]:.6f}, 差值={cx-initial_intrinsic[2]:.6f}")
    print(f"cy: 优化={cy:.6f}, 真值={initial_intrinsic[3]:.6f}, 差值={cy-initial_intrinsic[3]:.6f}")
    print(f"euler: 优化={euler}, 真值={initial_extrinsic[:3]}, 差值={euler-initial_extrinsic[:3]}")
    print(f"T_wc: 优化={T_wc_opt}, 真值={initial_extrinsic[3:]}, 差值={T_wc_opt -initial_extrinsic[3:]}")

    # ====== 计算并弹窗显示最终原点误差 ======
    # 用优化后的参数重新计算每个位置的原点误差
    Otw_list, _, _, _, _ = get_target_points(data)
    object_points = get_object_points()
    object_points_list = [object_points for _ in range(len(Otw_list))]
    # 用带有扰动的像素点
    image_points_list = perturbed_image_points_list
    origin_errors = []
    Otc_list = []
    Otc_estimated_list = []
    # 用优化参数计算R_wc, T_wc
    for Otw, img_pts, obj_pts in zip(Otw_list_noisy, image_points_list, object_points_list):
        K = np.array([[fx, 0, cx], [0, fy, cy], [0,  0,  1]])
        obj_pts = np.array(obj_pts, dtype=np.float32)
        img_pts = np.array(img_pts, dtype=np.float32)
        success, rvec, tvec = cv2.solvePnP(obj_pts, img_pts, K, None, flags=cv2.SOLVEPNP_EPNP)
        if not success:
            origin_errors.append([np.nan, np.nan, np.nan])
            Otc_list.append([np.nan, np.nan, np.nan])
            Otc_estimated_list.append([np.nan, np.nan, np.nan])
        else:
            R_tc, _ = cv2.Rodrigues(rvec)
            T_tc = tvec.flatten()
            Otc = R_wc_opt @ (Otw - T_cw_opt)
            Otc_estimated = R_tc @ np.zeros(3) + T_tc
            origin_error = Otc - Otc_estimated
            origin_errors.append(origin_error)
            Otc_list.append(Otc)
            Otc_estimated_list.append(Otc_estimated)
    final_cost = np.sum(np.array(origin_errors) ** 2)
    show_final_origin_errors(origin_errors, Otc_list, Otc_estimated_list, final_cost)

if __name__ == "__main__":
    main()
